import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as models
from rknn.api import RKNN


def export_pytorch_model():
    # 加载自定义模型结构
    model = models.mobilenet_v2(pretrained=False)

    # 修改分类器为三分类（健康、细菌感染、支原体感染）
    model.classifier = nn.Sequential(
        nn.Linear(1280, 256),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(256, 3),  # 修改为3个输出类别
    )

    # 加载训练好的权重
    model.load_state_dict(torch.load("./best_model.pth", map_location="cpu"))
    model.eval()

    # 创建示例输入（根据你的实际输入尺寸调整）
    example_input = torch.randn(1, 3, 224, 224)  # 保持与训练时相同的输入尺寸

    # 导出模型
    trace_model = torch.jit.trace(model, example_input)
    trace_model.save("./custom_model.pt")


def show_outputs(output):
    output_sorted = sorted(output, reverse=True)
    value = output_sorted[0]  # 只取最大值
    index = np.where(output == value)[0][0]

    # 定义类别映射
    class_names = {0: "健康", 1: "细菌性感染", 2: "支原体感染"}

    result_str = "\n----- 预测结果 -----\n"
    result_str += f"类别: {class_names[index]} (类别 {index})\n"
    result_str += f"概率: {value:.4f}\n"
    print(result_str)


def show_perfs(perfs):
    perfs = "perfs: {}\n".format(perfs)
    print(perfs)


def softmax(x):
    return np.exp(x) / sum(np.exp(x))


if __name__ == "__main__":
    export_pytorch_model()

    model = "./custom_model.pt"
    input_size_list = [[3, 224, 224]]

    # Create RKNN object
    rknn = RKNN()

    # pre-process config
    print("--> config model")
    rknn.config(
        mean_values=[[123.675, 116.28, 103.53]],
        std_values=[[58.395, 57.12, 57.375]],
        reorder_channel="0 1 2",
        target_platform="rv1126",
    )
    print("done")

    # Load pytorch model
    print("--> Loading model")
    ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
    if ret != 0:
        print("Load pytorch model failed!")
        exit(ret)
    print("done")

    # Build model
    print("--> Building model")
    ret = rknn.build(
        do_quantization=True,
        dataset="./dataset.txt",  # 需要准备自己的校准数据集
    )
    if ret != 0:
        print("Build pytorch failed!")
        exit(ret)
    print("done")

    # Export rknn model
    print("--> Export RKNN model")
    ret = rknn.export_rknn("./model.rknn")
    if ret != 0:
        print("Export model.rknn failed!")
        exit(ret)
    print("done")

    ret = rknn.load_rknn("./model.rknn")

    # Set inputs
    img = cv2.imread("./data/114.jpg")  # 使用自己的测试图像
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (224, 224))  # 根据模型输入尺寸调整
    # 如果需要，在此处添加与训练时相同的预处理（如归一化）

    # init runtime environment
    print("--> Init runtime environment")
    ret = rknn.init_runtime()
    if ret != 0:
        print("Init runtime environment failed")
        exit(ret)
    print("done")

    # Inference
    print("--> Running model")
    outputs = rknn.inference(inputs=[img])

    show_outputs(softmax(np.array(outputs[0][0])))
    print("done")

    # # perf
    # print('--> Begin evaluate model performance')
    # perf_results = rknn.eval_perf(inputs=[img])
    # print('done')

    rknn.release()
